# -*- coding: utf-8 -*-
"""
Created on Mon Feb  4 15:10:50 2013

@author: nzhao
"""

from traits.api import Int, List, Array
import scipy.sparse
import numpy as np

from phy.spin.MultiSpin import MultiSpin


class SparseHamiltonian(MultiSpin):
    dim = Int()
    col_list = Array()
    row_list = Array()
    val_list = Array()
    
    def __init__(self):
        self.col_list = np.array([])
        self.row_list = np.array([])
        self.val_list = np.array([])
        
class Many_SpinHalf_Hamiltonian(SparseHamiltonian):
    nspin = Int()
    basis_index = List()
    
    def __init__(self, n = 2):
        super(Many_SpinHalf_Hamiltonian, self).__init__()
        self.nspin = n
        self.dim   = 2**n
        self.basis_index = [bin(x)[2:].rjust(n,'0') for x in range(self.dim)]
    
    def get_mat(self):
        res = scipy.sparse.coo_matrix((self.val_list, (self.row_list, self.col_list)), shape=(self.dim, self.dim))
        return res
    
    def get_mat_dense(self):
        return self.get_mat().todense()
        
    def calc_elements(self, i, j):
        idx_i0_j0 = np.array([int(x,2) for x in self.basis_index if x[-i]=='0' and x[-j]=='0'])
        idx_i0_j1 = np.array([int(x,2) for x in self.basis_index if x[-i]=='0' and x[-j]=='1'])
        idx_i1_j0 = np.array([int(x,2) for x in self.basis_index if x[-i]=='1' and x[-j]=='0'])
        idx_i1_j1 = np.array([int(x,2) for x in self.basis_index if x[-i]=='1' and x[-j]=='1'])
        val00, val01, val10, val11 = self.elements(i, j)
        return [val00, val01, val10, val11]

    def elements(self, i, j):
        res00 = 1.0;
        res01 = 2.0;
        res10 = 3.0;
        res11 = 4.0;
        return [res00, res01, res10, res11]



if __name__=='__main__':
    hm = Many_SpinHalf_Hamiltonian(n=5)
    print hm.get_mat_dense()
    
    print hm.calc_elements(3,1)



















    